# boss_hp_detector_demo.py
import os
import cv2
import mss
import time
import numpy as np
import pydirectinput
from Activate_Window import activate_window
from config import (FPS, DURATION, CAPTURE_WINDOW, BOSS_HP_WINDOW)

OUTPUT_DIR = "./boss_hp_logs"
os.makedirs(OUTPUT_DIR, exist_ok=True)

# 原始画面截图
def capture_frame(sct):
    img = np.array(sct.grab(CAPTURE_WINDOW))    # BGRA (H, W, 4)
    return img[:, :, :3]                        # BGR  (H, W, 3)

# 从原图提取小骑士、BOSS血量图
def extract_hp_bar(frame,window):
    y = window["top"]
    x = window["left"]
    dy = window["height"]
    dx = window["width"]
    return frame[y:y+dy, x:x+dx]

# 识别boss血量图为数值信息
def detect_boss_hp_percentage(hp_bar_img):
    # [5,0,96]
    if not np.array_equal(hp_bar_img[0, 0], [5,0,96]):
        return 0.0
    w = hp_bar_img.shape[1]
    diff = np.diff(hp_bar_img[0, :, :].astype(np.int16), axis=0)
    candidates = np.where(np.abs(diff).max(axis=1) > 1)[0]
    if len(candidates) > 0:
        return candidates[0] / w
    else:
        return 1.0


def main():
    activate_window("Hollow Knight")

    pydirectinput.press("w")
    time.sleep(2)
    pydirectinput.press("w")
    time.sleep(1)
    pydirectinput.press("space")
    time.sleep(3)

    with mss.mss() as sct:
        start_time = time.time()
        count = 0
        prev_hp = 0

        while True:
            next_frame_time = start_time + (count + 1) / FPS
            current_time = time.time()
            if current_time < next_frame_time:
                time.sleep(next_frame_time - current_time)

            frame = capture_frame(sct)
            count += 1
            hp_bar = extract_hp_bar(frame,BOSS_HP_WINDOW)
            hp_percent = detect_boss_hp_percentage(hp_bar)

            # 血量图像仅用于调试，如果确定没有bug可以不保存图像只接收信号
            hp_str = f"{hp_percent:.2%}".replace('%', '')
            filename = os.path.join(OUTPUT_DIR, f"hp_{count:04d}_{hp_str}.png")
            cv2.imwrite(filename, hp_bar)

            '''
            if hp_percent == prev_hp:
                reward = 0.0
            elif prev_hp != 0 and hp_percent == 0:
                # 如果某次变化跨度过大，说明是已经被boss击败了
                # 这里输出的被击败的步数（时间点）是延后的（因为从角色死亡到boss血条消失有动画）
                # 不准确，仅仅是为了程序完整，判断真正死亡时间和反馈应该用小骑士血条。
                if (prev_hp - hp_percent) > 0.1:
                    reward = -5
                    print(f"第{count}步 : 已经被BOSS击败！奖励 {reward}")
                else:
                    reward = 5
                    print(f"第{count}步 : 击败Boss！奖励 {reward}")
                break
            elif hp_percent < prev_hp:
                reward = 1.0
                reward = (prev_hp - hp_percent)*10
                print(f"第{count}步 : 🩸 Boss掉血！奖励 {reward}")
            elif prev_hp == 0 and hp_percent != 0:
                reward = 5.0
                print(f"第{count}步 : 💥 初次攻击！奖励 {reward}")
            else:
                reward = -999
                print(f"第{count}步 : 血量异常增加 {reward}")
                break
            '''

            prev_hp = hp_percent

            if (time.time() - start_time) >= DURATION:
                print(f"{time.time() - start_time:.1f}s : ⏹️  采样结束（未击败 Boss）")
                break

        total_time = time.time() - start_time
        actual_fps = count / total_time  # 帧数 / 总时间

        print(f"\n📊 采样统计:")
        print(f"   目标时长: {DURATION:.1f}s")
        print(f"   实际耗时: {total_time:.2f}s")
        print(f"   总帧数: {count}")
        print(f"   实际平均采样率: {actual_fps:.2f} FPS")

if __name__ == '__main__':
    main()